import tensorflow as tf
import tensorflow as tf
import tf2onnx
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType
from onnxruntime.quantization import quantize_static, QuantType, QuantFormat, CalibrationMethod
            
def convert_to_tflite(pb_path, tflite_path,shape):
    converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
        graph_def_file=pb_path,
        input_arrays=["input_image"],  # 输入节点名称
        output_arrays=["output_image"],  # 输出节点名称
        input_shapes={"input_image": shape}  # 根据模型调整输入形状
    )
    
    # 优化（可选）
    # converter.optimizations = []
    # converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]    
    # converter.target_spec.supported_ops = []
    # 转换并保存
    tflite_model = converter.convert()
    with open(tflite_path, "wb") as f:
        f.write(tflite_model)

def tflite_to_onnx(tflite_path: str, onnx_path: str):
    """
    TFLite 转 ONNX
    :param tflite_path: TFLite 模型路径（如 "model.tflite"）
    :param onnx_path: 输出 ONNX 模型路径（如 "model.onnx"）
    """
    # 转换为 ONNX（需启用 TFLite 兼容模式）
    onnx_model, _ = tf2onnx.convert.from_tflite(
        tflite_path=tflite_path,
        opset=11,  # 推荐 ONNX 算子集版本（≥13）
        output_path=onnx_path
    )
    print(f"TFLite 转 ONNX 完成，保存至：{onnx_path}")
    return onnx_model
    #print(onnx_model)
    
    
# 2. 将 Div 节点替换为 Reciprocal + Mul
def replace_div_with_mul(onnx_model,path):
    graph = onnx_model.graph
    new_nodes = []
    
    for node in graph.node:
        graph = onnx_model.graph
    new_nodes = []
    
    for node in graph.node:
        # 处理 Div 节点
        if node.op_type == "Div":
            input_a = node.input[0]  # 被除数
            input_b = node.input[1]  # 除数
            output = node.output[0]
            
            # 创建 Reciprocal 节点 (计算 1/B)
            reciprocal_output = f"{input_b}_reciprocal"
            reciprocal_node = onnx.helper.make_node(
                "Reciprocal",
                inputs=[input_b],
                outputs=[reciprocal_output],
                name=f"{node.name}_reciprocal" if node.name else f"reciprocal_{input_b}"
            )
            new_nodes.append(reciprocal_node)
            
            # 创建 Mul 节点 (A * (1/B))
            mul_node = onnx.helper.make_node(
                "Mul",
                inputs=[input_a, reciprocal_output],
                outputs=[output],
                name=f"{node.name}_mul" if node.name else f"mul_{input_a}_{input_b}"
            )
            new_nodes.append(mul_node)
        
        
        else:
            new_nodes.append(node)
    
    # 替换图中的所有节点
    graph.ClearField("node")
    graph.node.extend(new_nodes)
    
    
    # 替换图中的所有节点
    graph.ClearField("node")
    graph.node.extend(new_nodes)
    onnx.save(onnx_model, path)    
    
convert_to_tflite("cartoon_bg.pb", "cartoon_bg.tflite",[320, 320, 3])
convert_to_tflite("cartoon_h.pb", "cartoon_h.tflite",[192, 192, 3])

model = tflite_to_onnx("cartoon_bg.tflite", "cartoon_bg_org.onnx")
replace_div_with_mul(model,"cartoon_bg.onnx")

model = tflite_to_onnx("cartoon_h.tflite", "cartoon_h.onnx")
# replace_div_with_mul(model,"cartoon_h.onnx")
